import os
import pickle
import time
from selenium.webdriver.remote.webdriver import WebDriver


class DriverCookie():
    """封装一个方法，以便以后调用"""

    def get_cookies(self, cookie_path: str, login_url: str, driver: WebDriver, callback,super_login=False):
        """获取cookies并保存到一个文件中"""
        print("请登录到 " + driver.title)
        driver.get(login_url)
        time.sleep(5)
        # 登录回调
        callback()

        cookies = {}
        # 有特殊登录地址就需要比较登录地址
        if super_login is False:
            # 是登录地址
            while driver.current_url == login_url:
                # 登录回调
                callback()
                cookies = driver.get_cookies()
                print('登陆后的cookies:')
                print(cookies)
                # driver.quit()
                if os.path.isfile(cookie_path) is False:
                    os.mkdir(os.path.basename(cookie_path))
                output_path = open(cookie_path, 'wb')  # 新建一个文件
                pickle.dump(cookies, output_path)
                output_path.close()
                # for cookie in cookies:
                #     cookies[cookie['name']] = cookie['value']
                #     output_path = open(cookie_path, 'wb')  # 新建一个文件
                #     pickle.dump(cookies, output_path)
                #     output_path.close()
        else:
            # 登录回调
            callback()
            cookies = driver.get_cookies()
            print('登陆后的cookies:')
            print(cookies)
            # driver.quit()
            output_path = open(cookie_path, 'wb')  # 新建一个文件
            pickle.dump(cookies, output_path)
            output_path.close()
            # for cookie in cookies:
            #     cookies[cookie['name']] = cookie['value']
            #     output_path = open(cookie_path, 'wb')  # 新建一个文件
            #     pickle.dump(cookies, output_path)
            #     output_path.close()
        # 刷新
        driver.get(login_url)
        return cookies

    # 读取cookie FIXME 登录 cookie 过期未自动处理
    def read_cookies(self, login_url: str, driver: WebDriver, callback):
        """这是读取cookies的操作
        :rtype: dict
        """

        # 如果有cookie就使用
        # 如果没有就登录，并重新设置cookie
        cookie_path = 'Browser\\cookies\\' + driver.name + '\\cookies.pickle'

        if os.path.exists(cookie_path):
            read_path = open(cookie_path, 'rb')
            cookies = pickle.load(read_path)
            print('读取Cookie:')
            print(cookies)
        else:
            print('Cookies不存在:登录获得cookies')
            cookies = self.get_cookies(cookie_path, login_url, driver, callback)
        # 设置浏览器cookies
        for cookie in cookies:
            driver.add_cookie(cookie)  # 加载Cookies
        return cookies
